import networkx as nx
import matplotlib.pyplot as plt
import numpy as np

# 这里不能用numpy的数组 权值会变成字符串 不知为何
# 节点可以是字符 权值是数字
paths=[
    ('A','B1',2),('B1','C1',3),('C1','D',1),('A','B2',4),('B2','C1',2),
    ('B2','C2',3),('B2','C3',1),('B1','C3',1),('B1','C2',3),('C2','D',3),
    ('C3','D',4)
]

G=nx.DiGraph()
G.add_weighted_edges_from(paths)
#nx.draw(G,with_labels=True)
#plt.show()
res=dict(nx.shortest_path_length(G,weight="weight"))
path=nx.shortest_path(G,'A','D',"weight")
print(res['A']['D'])
print(path)
# 用颜色区分最短路径
node_colors=[]
edge_colors=[]
for node in G.nodes():
    if(node in path):
        node_colors.append('red')
    else:
        node_colors.append('gray')
for edge in G.edges():
    if(set(edge).issubset(set(path))):
        edge_colors.append('red')
    else:
        edge_colors.append('gray')

#nx.draw(G,with_labels=True,node_color=node_colors,edge_color=edge_colors)
# 布局字典 即为每个节点的位置 注意画边也要传入这个pos
layout={
    'A':(20,50),'B1':(20,30),'B2':(10,70),'C1':(70,40),'C2':(60,40),'C3':(30,40),'D':(15,75),
}
nx.draw_networkx_nodes(G,pos=layout,node_colors=node_colors)
nx.draw_networkx_edges(G,pos=layout,width=2,edge_colors=edge_colors)
plt.show()